package goclient

import (
	"context"
	"errors"
	"github.com/gorilla/websocket"
)

var BufferNotAllSend = errors.New("buffer not all send")

type WebsocketConn struct {
	conn *websocket.Conn
}

func (c *WebsocketConn) Close() error {
	return c.conn.Close()
}

func (c *WebsocketConn) Send(msg *WebsocketMessage) error {
	return c.conn.WriteMessage(int(msg.MsgType), msg.Data)
}

func (c *WebsocketConn) Recv() (*WebsocketMessage, error) {
	msgType, data, err := c.conn.ReadMessage()
	if err != nil {
		return nil, err
	}
	return &WebsocketMessage{
		MsgType: int32(msgType),
		Data:    data,
	}, nil
}

func (c *WebsocketConn) Pipe(src ReadWriter) error {
	defer c.Close()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	errCh := make(chan error)

	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			default:
				msgType, data, err := c.conn.ReadMessage()
				if err != nil {
					errCh <- err
					return
				}
				if msgType == websocket.BinaryMessage || msgType == websocket.TextMessage {
					n, err := src.Write(data)
					if err != nil {
						errCh <- err
						return
					}
					if n != len(data) {
						errCh <- BufferNotAllSend
						return
					}
				}
				// TODO 可能要处理其他数据类型
			}
		}
	}()

	go func() {
		buffer := make([]byte, 2048)
		for {
			select {
			case <-ctx.Done():
				return
			default:
				n, err := src.Read(buffer)
				if err != nil {
					errCh <- err
					return
				}
				err = c.conn.WriteMessage(websocket.BinaryMessage, buffer[:n])
				if err != nil {
					errCh <- err
					return
				}
			}
		}
	}()

	return <-errCh
}
